import torch
import torch.nn as nn
import torch.nn.functional as F
from config import parsers


class PatchEmbedding(nn.Module):
    def __init__(self, img_size = 224, patch_size = 16, in_channels = 3, embed_dim = 768):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.in_channels = in_channels
        self.embed_dim = embed_dim
        self.num_patches = (img_size // patch_size) ** 2
        self.proj = nn.Conv2d(in_channels, embed_dim, kernel_size = patch_size, stride = patch_size)

    def forward(self, x):
        x = self.proj(x)  # (batch_size, embed_dim, num_patches ** 0.5, num_patches ** 0.5)
        x = x.flatten(2)
        x = x.transpose(-1, -2)
        return x


class BEiTImageEncoder(nn.Module):
    def __init__(self, img_size = 224, patch_size = 16, in_channels = 3, embed_dim = 768, num_layers = 12,
                 num_heads = 12, mlp_ratio = 4):
        super().__init__()
        self.patch_embed = PatchEmbedding(img_size = img_size, patch_size = patch_size, in_channels = in_channels,
                                          embed_dim = embed_dim)

        self.pos_embed = nn.Parameter(torch.rand(1, self.patch_embed.num_patches + 1, embed_dim))  # 图片块位置初始化编码
        # X = torch.arange(max_len, dtype = torch.float32).reshape(
        #     -1, 1) / torch.pow(10000, torch.arange(0, num_hiddens, 2, dtype = torch.float32) / num_hiddens)
        #
        # self.pos_embed[:, :, 0::2] = torch.sin(X)
        # self.pos_embed[:, :, 1::2] = torch.cos(X)

        self.cls_token = nn.Parameter(torch.rand(1, 1, embed_dim))  # cls初始化编码
        self.dropout = nn.Dropout(parsers().droput)
        self.fc = nn.Linear(768, parsers().num_classes)
        self.softmax = nn.Softmax(dim = 1)

        # Transformer Encoder
        self.transformer_encoder = nn.ModuleList()
        for _ in range(num_layers):
            self.transformer_encoder.append(
                nn.ModuleList([
                    nn.LayerNorm(embed_dim),  # 归一化
                    nn.MultiheadAttention(embed_dim, num_heads),
                    nn.Dropout(p = parsers().droput),
                    nn.LayerNorm(embed_dim),
                    # 连接层
                    nn.Sequential(nn.Linear(embed_dim, mlp_ratio * embed_dim),
                                  nn.GELU(),
                                  nn.Dropout(p = parsers().droput),
                                  nn.Linear(mlp_ratio * embed_dim, embed_dim),
                                  nn.Dropout(p = parsers().droput))
                ])
            )
        self.apply(self.init_weights)  # Module中的apply模块

    def init_weights(self, module):
        if isinstance(module, nn.Conv2d):
            nn.init.kaiming_normal_(module.weight, mode = 'fan_out')
            nn.init.constant_(module.bias, 0)
        elif isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, std = 0.02)
            nn.init.constant_(module.bias, 0)
        elif isinstance(module, nn.LayerNorm):
            nn.init.constant_(module.bias, 0)
            nn.init.constant_(module.weight, 1.0)

    def forward(self, x,h=None):
        x = self.patch_embed(x)
        if h is None:
            h=self.cls_token
        #input:[2,196,768]->output:[2,197,768]
        x = torch.cat((h.expand(x.shape[0], -1, -1), x), dim = 1)
        x = x + self.pos_embed  # [:, :x.shape[1], :]
        x = self.dropout(x)

        for layer_norm_1, attn, dropout_1, layer_norm_2, mlp in self.transformer_encoder:
            x_res = x
            x = layer_norm_1(x)
            x, _ = attn(x, x, x)
            x = dropout_1(x)
            # add残差连接
            x = x_res + x
            x_res = x
            x = layer_norm_2(x)
            x = mlp(x)
            x = dropout_1(x)
            x = x_res + x

        x = self.dropout(x)
        x = x[:, 0, :]  # 提取cls
        out = self.fc(x)
        # out = self.softmax(out)
        return out,h
